class GRO:
    def __init__(self, gro_name):
        with open(gro_name) as f:
            title = f.readline()
            atom_num = int(f.readline().strip())
            data = f.readlines()
            atom_list = [Atom(l) for l in data[:-1]]
            # 只保留边长, 忽略其他
            box_x, box_y, box_z = [float(t) for t in data[-1].split()[:3]]
        self.title = title
        self.atom_num = atom_num
        self.atom_list = atom_list
        self.box_x, self.box_y, self.box_z = box_x, box_y, box_z

    def __add__(self, in_gro):
        if not isinstance(in_gro, GRO):
            raise Exception('Error: Input not gro file!')
        self.atom_num += in_gro.get_atom_num()
        self.atom_list.extend(in_gro.get_atoms())
        return self

    def __iadd__(self, in_gro):
        if not isinstance(in_gro, GRO):
            raise Exception('Error: Input not gro file!')
        self.atom_num += in_gro.get_atom_num()
        self.atom_list.extend(in_gro.get_atoms())
        return self

    def to_gro(self, outname):
        with open(outname, 'w') as f:
            f.write(self.title)
            f.write(str(self.atom_num) + '\n')
            for atom in self.atom_list:
                f.write(str(atom) + '\n')
            f.write('{0:10.5f}{0:10.5f}{0:10.5f}\n'.format(self.box_x, self.box_y, self.box_z))

    def get_atoms(self):
        return self.atom_list

    def get_atom_num(self):
        return self.atom_num

    def __str__(self):
        return 'gro file: ' + self.title + 'containing ' + str(self.atom_num) + ' atoms in a ' + str(
            self.box_x) + ' × ' + str(self.box_y) + ' × ' + str(self.box_x) + ' box.'


class Atom:
    def __init__(self, at_line):
        self.resid, self.resname, self.atname, self.atid, self.x, self.y, self.z = \
            int(at_line[:5].strip()), at_line[5:11].strip(), at_line[11:15].strip(), int(at_line[15:20].strip()),\
                float(at_line[20:28].strip()), float(at_line[28:36].strip()), float(at_line[36:44].strip())
        if at_line[44:].strip() != '':
            print(at_line[44:].strip())
            self.x_speed, self.y_speed, self.z_speed = float(at_line[44:52].strip()), float(
                at_line[52:60].strip()), float(at_line[60:68].strip())
        else:
            self.x_speed, self.y_speed, self.z_speed = None, None, None

    def __expr__(self):
        if self.x_speed is not None:
            return '{0:5d}{1:6s}{2:>4s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}{7:8.3f}{8:8.3f}{9:8.3f}' \
                .format(self.resid, self.resname, self.atname, self.atid, self.x, self.y, self.z, self.x_speed,
                        self.y_speed, self.z_speed)
        return '{0:5d}{1:6s}{2:>4s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}' \
            .format(self.resid, self.resname, self.atname, self.atid, self.x, self.y, self.z)

    def __str__(self):
        if self.x_speed is not None:
            return '{0:5d}{1:6s}{2:>4s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}{7:8.3f}{8:8.3f}{9:8.3f}' \
                .format(self.resid, self.resname, self.atname, self.atid, self.x, self.y, self.z, self.x_speed,
                        self.y_speed, self.z_speed)
        return '{0:5d}{1:6s}{2:>4s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}' \
            .format(self.resid, self.resname, self.atname, self.atid, self.x, self.y, self.z)


if __name__ == '__main__':
    conf = GRO('test/combine_gro/conf.gro')
    SMP = GRO('test/combine_gro/SMP_18col.gro')
    conf += SMP
    conf.to_gro('test/combine_gro/combine.gro')
